import sys
import os
import json
import logging
import numpy as np
import torch

from dataclasses import dataclass

from ascendebug.features.dump.dump_accuracy_checkpoint import set_acc_dump_dir

logging.basicConfig(filename='gen_bin_golden.log', level=logging.DEBUG)

DTYPE = {"float16": np.float16, "float32": np.float32, "int16": np.int16, "int32": np.int32, "int64": np.int64}


def get_input_attrs_content(input_json: str) -> dict:
    content = {}
    input_json = os.path.realpath(input_json)
    with open(input_json, "r") as in_hadle:
        try:
            content = json.load(in_hadle)
        except json.decoder.JSONDecodeError:
            logging.error("json.decoder.JSONDecodeError")
        finally:
            pass
    return content


def get_data_file(res_path, contents):
    data_files = {}
    for content in contents:
        if content["data_file"] is None:
            continue
        data_files[content["name"]] = os.path.join(res_path, content["data_file"])
    return data_files


@dataclass
class OperatorInfo:
    xyz1_shape: tuple
    xyz2_shape: tuple


def get_args():
    """
    需要用户按需配置
    sys.argv[1]: input data file path
    sys.argv[2]: output file path
    sys.argv[3]: gen data type: input/output/all
    """
    logging.info(f"data file path: {sys.argv[2]}")
    input_file = sys.argv[1]
    res_path = sys.argv[2]
    if len(sys.argv) > 3:
        checkpoint_dump_path = sys.argv[3]
        set_acc_dump_dir(checkpoint_dump_path)
    content = get_input_attrs_content(input_file)
    res_abs_path = os.path.abspath(res_path)
    if not os.path.exists(res_abs_path):
        os.popen(f'mkdir -p {res_abs_path}')

    in_content = content.get("inputs")
    out_content = content.get("outputs")
    input_files = get_data_file(res_path, in_content)
    output_files = get_data_file(res_path, out_content)

    xyz1_shape = in_content[0].get("shape")
    xyz2_shape = in_content[1].get("shape")

    op_info = OperatorInfo(xyz1_shape, xyz2_shape)
    return op_info, input_files, output_files


def run_three_nn(b, n, m, xyz1, xyz2):
    dist = np.zeros_like(xyz1, dtype=np.float32)
    idx = np.zeros_like(xyz1, dtype=np.int32)

    for ii in range(b):
        for jj in range(n):
            x1 = xyz1[ii, jj, 0]
            y1 = xyz1[ii, jj, 1]
            z1 = xyz1[ii, jj, 2]
            best1 = 1e40
            best2 = 1e40
            best3 = 1e40
            besti1 = 0
            besti2 = 0
            besti3 = 0
            for kk in range(m):
                x2 = xyz2[ii, kk, 0]
                y2 = xyz2[ii, kk, 1]
                z2 = xyz2[ii, kk, 2]
                d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1)
                if (d < best1):
                    best3 = best2
                    besti3 = besti2
                    best2 = best1
                    besti2 = besti1
                    best1 = d
                    besti1 = kk
                elif (d < best2):
                    best3 = best2
                    besti3 = besti2
                    best2 = d
                    besti2 = kk
                elif (d < best3):
                    best3 = d
                    besti3 = kk
            dist[ii, jj, 0] = best1
            idx[ii, jj, 0] = besti1
            dist[ii, jj, 1] = best2
            idx[ii, jj, 1] = besti2
            dist[ii, jj, 2] = best3
            idx[ii, jj, 2] = besti3

    return dist, idx


def gen_data_and_golden(op_info, input_paths, golden_paths):
    logging.info(f"op_info={op_info}")
    xyz1 = np.random.uniform(-10, 10, op_info.xyz1_shape).astype(np.float32)
    xyz1.tofile(input_paths["xyz1"])
    logging.info(f"xyz1={xyz1}")
    xyz2 = np.random.uniform(-10, 10, op_info.xyz2_shape).astype(np.float32)
    xyz2.tofile(input_paths["xyz2"])
    logging.info(f"xyz2={xyz2}")

    golen_dist, golden_idx = run_three_nn(
        op_info.xyz1_shape[0], op_info.xyz1_shape[1], op_info.xyz2_shape[1], xyz1, xyz2
    )

    logging.info(f"golden_y={golen_dist}")
    golen_dist.tofile(golden_paths["dist"])
    logging.info(f"golden_y={golden_idx}")
    golden_idx.tofile(golden_paths["indices"])


if __name__ == '__main__':
    """
    sys.argv[1]: input data file path
    sys.argv[2]: output file path
    sys.argv[3]: gen data type: input/output/all
    """
    gen_data_and_golden(*get_args())
